
import os
import numpy as np
import cv2
from shutil import rmtree


def empty_dir(path):
    """ empty specified dir """
    if os.path.exists(path):
        rmtree(path)
    os.mkdir(path)


if __name__ == "__main__":
    img_dir = r"E:\xianhuang0818\xhsi_seg20210113\images"
    equ_dir = r"E:\xianhuang0818\xhsi_seg20210113\img_equ"
    mask_dir = r"E:\xianhuang0818\xhsi_seg20210113\masks"
    roi_dir = r"E:\xianhuang0818\xhsi_seg20210113\masks_roi"
    ori_dir = r"E:\xianhuang0818\xhsi_seg20210113\images_ori"
    empty_dir(roi_dir)
    empty_dir(ori_dir)
    empty_dir(equ_dir)
    names = os.listdir(mask_dir)
    for name in names:
        print(name)
        img_path = os.path.join(img_dir, name)
        equ_path = os.path.join(equ_dir, name)
        mask_path = os.path.join(mask_dir, name)
        roi_path = os.path.join(roi_dir, name)
        ori_path = os.path.join(ori_dir, name)
        img = cv2.imread(img_path)
        cv2.imwrite(ori_path, img)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img_equ = cv2.equalizeHist(gray)
        cv2.imwrite(equ_path, img_equ)
        img_mask = cv2.imread(mask_path, 0)
        ori_size = img.shape
        imgShow = np.empty((ori_size[0], ori_size[1], 3), np.uint8)
        imgShow[:, :, 0] = img[:, :, 0]
        imgShow[:, :, 1] = img[:, :, 1]
        imgShow[:, :, 2] = np.maximum(img[:, :, 2], img_mask) * 0.8 + img[:, :, 2] * 0.2
        font = cv2.FONT_HERSHEY_SIMPLEX
        imgShow = cv2.putText(imgShow, name, (20, 100), font, 1.2, (255, 255, 255), 2)
        cv2.imshow("imgShow", cv2.resize(imgShow, (900, 900)))

        imgShow2 = np.empty((ori_size[0], ori_size[1], 3), np.uint8)
        imgShow2[:, :, 0] = img_equ
        imgShow2[:, :, 1] = img_equ
        imgShow2[:, :, 2] = np.maximum(img_equ, img_mask) * 0.8 + img_equ * 0.2
        imgShow2 = cv2.putText(imgShow2, name, (20, 100), font, 1.2, (255, 255, 255), 2)
        cv2.imshow("imgShow2", cv2.resize(imgShow2, (900, 900)))

        cv2.waitKey(1)
        cv2.imwrite(roi_path, imgShow)




